GDN prefill reset currently selects cache slots with a CUDA boolean mask before assigning zeros. The boolean-indexed selection has a data-dependent output shape, which forces the host to obtain the selected element count and introduces a stream synchronization for every GDN layer.

Replace the two indexed assignments with a fixed-shape Triton launch. Each program reads the request's cache index and initialization flag on device, validates the cache slot, and clears both the SSM and convolution state rows in one kernel. Requests with initialized state or invalid sentinel indices remain untouched.

Add a CUDA unit test covering reset slots, preserved initialized slots, untouched unrelated slots, and the negative invalid-index sentinel.

Dev Engineer Review

• Added a Triton kernel helper (_reset_gdn_states_kernel / _reset_gdn_states) in tensorrt_llm/_torch/modules/mamba/gdn_mixer.py to reset eligible GDN state rows (SSM and convolution) for a batch of requests by using state_indices and has_initial_states directly on device.
• Kernel behavior:
  • Loads state_idx and computes needs_reset = ~has_initial_states[request].
  • Validates state_idx via 0 <= state_idx < NUM_CACHE_LINES (treats -1 as invalid).
  • Zeroes ssm_states and conv_states only when needs_reset & valid_state, with additional bounds offsets < SSM_STATE_SIZE / offsets < CONV_STATE_SIZE (row-size gating).
  • Leaves initialized states unchanged; invalid sentinel indices are not zeroed.
• Updated forward_core to replace the prior masked/host-influenced reset logic for prefills with a single call:
  • _reset_gdn_states(ssm_states, conv_states, state_indices_p, has_initial_states_p) when num_prefills > 0.
• Review focus: correctness of Triton grid/shape assumptions (grid = (num_requests, cdiv(max(state_sizes), 256))), index/sentinel handling, and ensuring the masked store prevents any out-of-bounds effects even when state_idx is invalid.

QA Engineer Review

• Added CUDA-only unit test test_reset_gdn_states_preserves_initialized_and_invalid_slots in tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py.
  • Constructs a state_pool backing views for ssm_states/conv_states.
  • Calls _reset_gdn_states with state_indices=[1, 3, -1] and has_initial_states=[False, True, False].
  • Asserts:
    • Slot 1 is zeroed.
    • Initialized slot 3 is preserved.
    • Invalid -1 does not modify any unrelated slot(s).
• Test coverage in tests/integration/test_lists/ (test-db/qa) for this specific unit test could not be verified from search results.
• Verdict: needs follow-up.

Description

Test Coverage

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.